import cv2
import numpy as np
import base64
from paddlers.tasks.utils.visualize import visualize_detection
import time

from skimage.io import imread, imsave
from PIL import Image
from matplotlib import pyplot as plt
import paddle

'''
变化检测后处理函数
@prob_map: 网络输出的掩码图
@img_a: 前时相图
@img_b: 后时相图
'''
def CDpostProcess(prob_map, img_a=None):
    #后处理
    begin = time.time()
    # print(img_a.shape)
    # cm_slide = (prob_map>0.5).astype('int32')
    end = time.time()
    mask=prob_map.astype(np.uint8)
    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(img_a, contours, -1, (179, 179, 255), -1)

    img_a = img_a[:, :, ::-1]
    img_a[..., 2] = np.where(mask == 1, 51, img_a[..., 2])

    return img_a, end-begin

def onlyshowOne(prob_map, image, param):
    #地物分类仅显示一类的后处理函数
    begin = time.time()
    if(param!=-1):
        for i,k in enumerate(prob_map):
            for j,c in enumerate(k):
                if (c!=param):
                    c = 0
    
    print(image.shape,prob_map.shape,param)
    cm_slide = (prob_map>0.5).astype('int32')
    end = time.time()
    return cm_slide*255, end-begin

def PDetProcess(prob_map, img_a=None, param=-1):
    #后处理
    #将单通道label映射成三通道图像，cv2.COLORMAP_RAINBOW调整映射后的颜色
    begin = time.time()
    if(int(param)!=-1):
        for i,k in enumerate(prob_map):
            for j,c in enumerate(k):
                if (c!=int(param)):
                    prob_map[i,j] = 4
    prob_map = cv2.applyColorMap((prob_map*(255/4)).astype(np.uint8), cv2.COLORMAP_RAINBOW)
    
    for i,k in enumerate(prob_map):
        for j,c in enumerate(k):
            #print(c) #查background的颜色，改成[0，0，0]
            if not (c - [255,0,170]).any():
                prob_map[i,j,:] = [0,0,0]
    rows, cols = prob_map.shape[:2]
    roi = img_a[:rows,:cols]
    #print(prob_map.shape)
    labgray = cv2.cvtColor(prob_map, cv2.COLOR_BGR2GRAY)
    _, mask = cv2.threshold(labgray, 10, 255, cv2.THRESH_BINARY)
    mask_inv = cv2.bitwise_not(mask)

    img_bg = cv2.bitwise_and(roi,roi,mask=mask_inv)
    res = cv2.add(img_bg,prob_map)
    end = time.time()
    return res, end-begin

'''
目标检测后处理函数
@prob_map: 网络输出的标签
@img_a: 输入原图
@threshold: 置信度阈值
'''

def DetpostProcess(prob_map, img_a, threshold=0.2):
    #后处理
    #print(prob_map.shape)
    #vis_res = []
    begin = time.time()
    if len(prob_map) > 0:
        img_a = visualize_detection(
            np.array(img_a), prob_map,
            color=np.asarray([[255,0,0]], dtype=np.uint8),
            threshold=float(threshold), save_dir=None
        )
    #cm_slide = (prob_map>0.5).astype('int32')
    #vis_res.append(img_a)
    #print(type(img_a))
    end = time.time()
    return np.array(img_a), end-begin


'''
目标提取后处理函数
'''

def SegpostProcess(prob_map,img_a):
    #后处理
    begin = time.time()
    # print(img_a.shape)
    #cm_slide = (prob_map>0.5).astype('int32')
    mask=prob_map.astype(np.uint8)
    
    
    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(img_a, contours, -1, (179, 179, 255), -1)

    img_a = img_a[:, :, ::-1]
    img_a[..., 2] = np.where(mask == 1, 255, img_a[..., 2])
    end = time.time()
    return np.array(img_a), end-begin 

'''
opencv格式和base64编码相互转换
'''
def base64_to_image(img):
    if(len(img)%3 == 1): 
        img += "=="
    elif(len(img)%3 == 2): 
        img += "=" 
    img = base64.b64decode(img)
    #print(img)
    img = np.fromstring(img, np.uint8) # 转换np序列
    #print(img.shape)
    img = cv2.imdecode(img,cv2.COLOR_BGR2RGB)  # 转换Opencv格式
    #img = cv2.imdecode(img)
    #print("107:",img)
    return img

def image_to_base64(image):
    #print(image.shape)
    data = cv2.imencode('.png', image)[1]
    return base64.b64encode(data.tostring()).decode('utf8')

'''
用于将目标检测的标签从字符串格式转化为可用的list格式
'''
def Str_to_Detlabel(labelstr):
    labellist = labelstr.strip("[").strip("]").split("}")[:-1]
    label = []
    for onelabel in labellist:
        i1,i2,i3,i4 = onelabel.index("\'category_id\'"), onelabel.index("\'category\'"), onelabel.index("\'bbox\'"), onelabel.index("\'score\'")
        label.append({
            'category_id':int(onelabel[i1+len("\'category_id\'")+2:i2-2]),
            'category':onelabel[i2+len("\'category\'")+2:i3-2].strip("\'"),
            'bbox':[float(x) for x in onelabel[i3+len("\'bbox\'")+2:i4-2].strip("[").strip("]").split(',')],
            'score':float(onelabel[i4+len("\'score\'")+2::])
        })
    return label

'''
将切片还原成整图
@patches: 图块推理结果
@ori_size: 原图尺寸
@window_size: 窗口大小
@stride: 
'''
def recons_prob_map(patches, ori_size, window_size, stride):
    """从裁块结果重建原始尺寸影像"""
    h, w = ori_size
    win_gen = WindowGenerator(h, w, window_size, window_size, stride, stride)
    prob_map = np.zeros((h,w), dtype=np.float)
    cnt = np.zeros((h,w), dtype=np.float)
    # XXX: 需要保证win_gen与patches具有相同长度。此处未做检查
    for (rows, cols), patch in zip(win_gen, patches):
        prob_map[rows, cols] += patch
        cnt[rows, cols] += 1
    prob_map /= cnt
    return prob_map

'''
用于滑窗推理
'''
class WindowGenerator:
    def __init__(self, h, w, ch, cw, si=1, sj=1):
        self.h = h
        self.w = w
        self.ch = ch
        self.cw = cw
        if self.h < self.ch or self.w < self.cw:
            raise NotImplementedError
        self.si = si
        self.sj = sj
        self._i, self._j = 0, 0

    def __next__(self):
        # 列优先移动（C-order）
        if self._i > self.h:
            raise StopIteration
        
        bottom = min(self._i+self.ch, self.h)
        right = min(self._j+self.cw, self.w)
        top = max(0, bottom-self.ch)
        left = max(0, right-self.cw)

        if self._j >= self.w-self.cw:
            if self._i >= self.h-self.ch:
                # 设置一个非法值，使得迭代可以early stop
                self._i = self.h+1
            self._goto_next_row()
        else:
            self._j += self.sj
            if self._j > self.w:
                self._goto_next_row()

        return slice(top, bottom, 1), slice(left, right, 1)

    def __iter__(self):
        return self

    def _goto_next_row(self):
        self._i += self.si
        self._j = 0
